using System;
using System.Collections.Generic;
using System.Text;
using System.Collections;
using System.Collections.Specialized;

namespace Brett
{
    public class VirusDatabase
    {
        Db4objects.Db4o.IObjectContainer _database = null;

        public VirusDatabase(string database)
        {
            if (!System.IO.File.Exists(database))
            {
                try
                {
                    System.IO.File.Create(database).Close();
                }
                catch
                {
                    throw;
                }
            }
            try
            {
                _database = Db4objects.Db4o.Db4oFactory.OpenFile(database);
            }
            catch
            {
                throw;
            }
        }

        public bool AddVirus(string file, string name)
        {
            byte[] fileBytes = null;
            System.IO.MemoryStream ms = new System.IO.MemoryStream();
            try
            {

                byte[] bytes = System.IO.File.ReadAllBytes(file);

                for (int i = 0; i < bytes.Length; i++)
                {
                    if (!bytes[i].Equals(0))
                    {
                        ms.WriteByte(bytes[i]);
                    }
                }
                ms.Flush();
                fileBytes = ms.ToArray();
                ms.Close();
            }
            catch
            {
                return false;
            }

            //Split the file up for hashing
            int splitCount = 10;
            
            if (fileBytes.Length >= splitCount)
            {
                int splitLength = fileBytes.Length / splitCount;

                using (System.Security.Cryptography.MD5CryptoServiceProvider md5 = new System.Security.Cryptography.MD5CryptoServiceProvider())
                {
                    StringCollection hashes = new StringCollection();
                    hashes.Add(GetHashString(md5.ComputeHash(fileBytes, 0, splitLength)));
                    for (int i = 1; i < splitCount; i++)
                    {
                        hashes.Add(GetHashString(md5.ComputeHash(fileBytes, (splitLength * i), splitLength)));
                    }

                    Virus virus = new Virus(name);
                    if (_database.Get(new Virus(name)).Count > 0)
                    {
                        return false;
                    }
                    virus.ChunkLength = splitLength;

                    virus.Hashes = StringCollectionToStringArray(hashes);
                    hashes.Clear();

                    _database.Set(virus);
                    _database.Commit();
                    return true;
                }
            }
            return false;
        }

        public string ScanForVirus(string file)
        {
            byte[] fileBytes = null;
            System.IO.MemoryStream fileByteStream = new System.IO.MemoryStream();
            
            byte[] bytes = ForceFileRead(file);
            if (bytes == null)
            {
                return string.Empty;
            }

            for (int i = 0; i < bytes.Length; i++)
            {
                if (!bytes[i].Equals(0))
                {
                    fileByteStream.WriteByte(bytes[i]);
                }
            }
            fileBytes = fileByteStream.ToArray();
            fileByteStream.Close();

            Db4objects.Db4o.IObjectSet query = _database.Get(null);
            foreach (Virus virus in query)
            {
                if (fileBytes.Length >= virus.ChunkLength)
                {
                    byte[] temp = new byte[virus.ChunkLength];
                    System.IO.MemoryStream ms = new System.IO.MemoryStream(fileBytes);
                    System.Security.Cryptography.MD5CryptoServiceProvider md5 = new System.Security.Cryptography.MD5CryptoServiceProvider();
                    int read = 0;
                    string hash = string.Empty;
                    for (int i = 0; i < ms.Length; i += read)
                    {
                        read = ms.Read(temp, 0, temp.Length);
                        hash = GetHashString(md5.ComputeHash(temp));
                        foreach (string s in virus.Hashes)
                        {
                            if (hash.Equals(s))
                            {
                                ms.Close();
                                md5.Clear();
                                return virus.Name;
                            }
                        }
                        if (read < virus.ChunkLength)
                        {
                            break;
                        }
                    }
                }
            }
            return string.Empty;
        }

        public bool Close()
        {
            return _database.Close();
        }

        public IList<Brett.VirusDatabase.Virus> GetAllViruses()
        {
            return _database.Query<Brett.VirusDatabase.Virus>();
        }

        private string GetHashString(byte[] hash)
        {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < hash.Length; i++)
            {
                sb.Append(hash[i].ToString("x2"));
            }
            return sb.ToString();
        }

        private string[] StringCollectionToStringArray(StringCollection sc)
        {
            string[] temp = new string[sc.Count];
            for (int i = 0; i < temp.Length; i++)
            {
                temp[i] = sc[i];
            }
            return temp;
        }

        private byte[] ForceFileRead(string file)
        {
            try
            {
                System.IO.FileStream fs = new System.IO.FileStream(file, System.IO.FileMode.Open, System.IO.FileAccess.Read, System.IO.FileShare.None);
                System.IO.MemoryStream ms = new System.IO.MemoryStream();
                byte[] buffer = new byte[1000];
                int read;
                while (true)
                {
                    read = fs.Read(buffer, 0, buffer.Length);
                    ms.Write(buffer, 0, read);
                    if (read < buffer.Length)
                    {
                        break;
                    }
                }
                fs.Close();
                return ms.ToArray();
            }
            catch
            {
                return null;
            }

            
        }

        public class Virus
        {
            string _name;
            int _chunkLength;
            string[] _hashes;

            public Virus(string name)
            {
                _name = name;
            }

            public string[] Hashes
            {
                get
                {
                    return _hashes;
                }
                set
                {
                    _hashes = value;
                }
            }

            public string Name
            {
                get
                {
                    return _name;
                }
            }

            public int ChunkLength
            {
                get
                {
                    return _chunkLength;
                }
                set
                {
                    _chunkLength = value;
                }
            }
        }


    }

    public static class Test
    {
        public static void Tester()
        {
            VirusDatabase vd = new VirusDatabase(@"C:\Documents and Settings\Brett Barnett\Desktop\viruses.db");
            vd.AddVirus(@"C:\Documents and Settings\Brett Barnett\Desktop\PI2.3.2\piserver.exe", "Trojan.Chasenet.PoisonIvy3.2");
            string result = vd.ScanForVirus(@"C:\Documents and Settings\Brett Barnett\Desktop\PI2.3.2\piserver2.exe");
            vd.Close();
        }
    }
}
